
import logging
logging.addLevelName(logging.WARN,  'WARN')
logging.addLevelName(logging.FATAL, 'FATAL')
logging.basicConfig(level=logging.DEBUG, 
                    format='%(asctime)s %(levelname)-5s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    force=True)
import sys, safe_sys
safe_sys.stderr  = sys.stderr
safe_sys.stdout  = sys.stdout

import builtins
import itertools
import lobobt.formula as formula
import lobobt.strategy
import random

def get_safe_builtins(*ext_safe_names:str):    
    _safe_names = [
        '__build_class__',
        'None',
        'False',
        'True',
        'abs',
        'bool',
        'bytes',
        'callable',
        'chr',
        'complex',
        'divmod',
        'float',
        'hash',
        'hex',
        'id',
        'int',
        'isinstance',
        'issubclass',
        'len',
        'oct',
        'ord',
        'pow',
        'range',
        'repr',
        'round',
        'slice',
        'sorted',
        'str',
        'tuple',
        'zip',
    ]

    _safe_exceptions = [
        'ArithmeticError',
        'AssertionError',
        'AttributeError',
        'BaseException',
        'BufferError',
        'BytesWarning',
        'DeprecationWarning',
        'EOFError',
        'EnvironmentError',
        'Exception',
        'FloatingPointError',
        'FutureWarning',
        'GeneratorExit',
        'IOError',
        'ImportError',
        'ImportWarning',
        'IndentationError',
        'IndexError',
        'KeyError',
        'KeyboardInterrupt',
        'LookupError',
        'MemoryError',
        'NameError',
        'NotImplementedError',
        'OSError',
        'OverflowError',
        'PendingDeprecationWarning',
        'ReferenceError',
        'RuntimeError',
        'RuntimeWarning',
        'StopIteration',
        'SyntaxError',
        'SyntaxWarning',
        'SystemError',
        'SystemExit',
        'TabError',
        'TypeError',
        'UnboundLocalError',
        'UnicodeDecodeError',
        'UnicodeEncodeError',
        'UnicodeError',
        'UnicodeTranslateError',
        'UnicodeWarning',
        'UserWarning',
        'ValueError',
        'Warning',
        'ZeroDivisionError',
    ]

    return {name: getattr(builtins, name) 
            for name in itertools.chain(_safe_names, _safe_exceptions, ext_safe_names)}

if __name__ == '__main__':
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--data', '-d', type=str, default=None, help='data file path')
    parser.add_argument(dest='script', type=str, help='script file path')
    # parser.add_argument(dest='args', type=str, nargs='*', help='arguments for script')

    args = parser.parse_args()

    import pandas as pd
    
    if args.data is not None:
        datafile:str = args.data
        if datafile.endswith('.csv'):
            df = pd.read_csv(datafile, sep=None, engine='python')
        elif datafile.endswith('.tsv'):
            df = pd.read_csv(datafile, sep='\t')
        elif datafile.endswith('.xlsx'):
            df = pd.read_excel(datafile)
        else:
            raise ValueError('data file must be csv or xlsx')
        datas = df.to_dict('records')
        print(df)
        for x in datas:
            print(x)
    else:
        datas = [None]

    with open(args.script, 'r') as f:
        script = f.read()

    safe_builtins = get_safe_builtins('print')

    safe_imports = ['math', 'numpy', 'pandas', 'polars', 'datetime', 'time', 'random', 
                    'lobobt.formula', 'lobobt.strategy',
                    ]
    alt_imports = {'sys': safe_sys, 
                   'lobobt.formula': formula, 
                   'lobobt.strategy': lobobt.strategy
                   }

    def safe_import(name, *args, **kwargs):
        if name in alt_imports:
            return alt_imports[name]
        if name in safe_imports:
            return __import__(name)
        else:
            raise ImportError(f'not allowed to import: {name}')
        
    safe_builtins['__import__'] = safe_import
    
    global_context = {
        '__builtins__': safe_builtins,
        '__name__': '__main__',
    }

    bcode = compile(script, '<script>', 'exec',)

    formula.logger = logging.getLogger('lobobt')
    
    for i, data in enumerate(datas):
        formula._datas = data
        formula.PERIOD = i
        d = formula.__dict__
        for name in ['CODE', 'DATE', 'OPEN', 'HIGH', 'LOW', 'CLOSE', 'VOLUME', 'AMOUNT']:
            if data is not None:
                if name in data:
                    d[name] = data[name]
            else:
                if name == 'DATE':
                    d['DATE'] = '2025-05-06'
                elif name == 'CODE':
                    d['CODE'] = '000001.sz'
                else:
                    d[name] = random.randint(1, 100)
        exec(bcode, global_context)
